from langchain.chat_models import ChatOpenAI
from langchain.prompts import ChatPromptTemplate
from langchain.schema import StrOutputParser
from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
from langchain.schema.runnable import RunnablePassthrough
from pydantic import BaseModel,Field,validator
from typing import List,Dict,Optional
from enum import Enum
from dotenv import load_dotenv,find_dotenv
_=load_dotenv(find_dotenv())


# 输出结构
class SortEnum(str,Enum):
  data  = 'data'
  price = 'price'

class OrderIngEnum(str,Enum):
  ascend= 'ascend'
  descend= 'descend'

class Semantics(BaseModel):
  name:Optional[str] = Field(default=None,description="流量包名称")
  price_lower :Optional[int] = Field(default=None,description="流量包价格下限")
  price_upper :Optional[int] = Field(default=None,description="流量包价格上限")
  sort:Optional[SortEnum] = Field(default=None,description="按价格或流量排序")
  data_lower :Optional[int] = Field(default=None,description="流量下限")
  data_upper :Optional[int] = Field(default=None,description="流量上限")
  ordering:Optional[OrderIngEnum] = Field(default=None,description="升序或降序")


#outputParser
parser = PydanticOutputParser(pydantic_object=Semantics) 

"""
 ChatPromptTemplate 类来创建一个聊天模板
from_messages() 方法接受一个消息列表作为输入，每个消息都由一个角色和一个内容组成。

partial() 方法可以用于创建一个新的函数，该函数将原始函数的某些参数固定为特定值。在本例中，partial() 方法
将 format_instructions 参数固定为 parser.get_format_instructions() 的返回值。
parser.get_format_instructions() 是一个函数，它返回一个字符串，该字符串描述了输出 JSON 的格式。
"""
prompt = ChatPromptTemplate.from_messages([
  (
    "system",
    "将用户的输入解析为json表示，输出格式如下:\n{format_instructions}\n不要输入未提及的字段"
  ),
  ('human',"{query}")
]).partial(format_instructions=parser.get_format_instructions())

# 模型
model = ChatOpenAI(temperature=0)


# chain
runable = (
  {'query':RunnablePassthrough()} | prompt | model | parser 
)

# 运行
print(runable.invoke("不超过100元的流量大的套餐有哪些"))